import numpy as np
import mcubes
import os
import argparse
from PyCmpltrtok.common import rand_name_on_now

parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('-N', help='N', type=int, default=128)
parser.add_argument('--nonsparse', help='if nonsparse', action='store_true')
parser.add_argument('--ij', help='if indexing ij', action='store_true')
args = parser.parse_args()

scale = 1.0
N = args.N
sparse = not args.nonsparse
ij = args.ij
if ij:
    ij = 'ij'
else:
    ij = 'xy'

print('scale', scale)
print('N', N)
print('sparse', sparse)
print('indexing', ij)

tx = np.linspace(-scale, scale, N)
ty = np.linspace(-scale, scale, N)
tz = np.linspace(-scale, scale, N)

# N = 512 sparse RAM 2.28G
# N = 512 not sparse RAM 6.24G
X, Y, Z = np.meshgrid(tx, ty, tz, sparse=sparse, indexing=ij)


def get_sphere(X, Y, Z, x, y, z, r):
    u = (X - x) ** 2 + (Y - y) ** 2 + (Z - z) ** 2 < r ** 2
    return u


sphere0 = get_sphere(X, Y, Z, -scale, -scale, -scale, 0.05)
spherex = get_sphere(X, Y, Z, -scale + 0.1, -scale, -scale, 0.05)
spherey1 = get_sphere(X, Y, Z, -scale, -scale + 0.1, -scale, 0.05)
spherey2 = get_sphere(X, Y, Z, -scale, -scale + 0.2, -scale, 0.05)
spherez1 = get_sphere(X, Y, Z, -scale, -scale, -scale + 0.1, 0.05)
spherez2 = get_sphere(X, Y, Z, -scale, -scale, -scale + 0.2, 0.05)
spherez3 = get_sphere(X, Y, Z, -scale, -scale, -scale + 0.3, 0.05)

volume = sphere0 + spherex + spherey1 + spherey2 + spherez1 + spherez2 + spherez3
print('volume', volume.shape)

volume[0:1, ::4, ::4] = 1000.  # x: yz
volume[::8, 0:1, ::8] = 1000.  # y: xz
volume[::16, ::16, 0:2] = 1000.  # z: xy

vertices, triangles = mcubes.marching_cubes(volume, 0)
print('vertices', vertices.shape)
print('triangles', triangles.shape)
xfilename = os.path.splitext(os.path.split(__file__)[1])[0]
path = f'_save/{N}_{ij}_{xfilename}.{rand_name_on_now()}.dae'
print(f'Saving to {path}')
mcubes.export_mesh(vertices, triangles, path, 'MyXYZ')
print('Saved')
path = f'_save/{N}_{ij}_{xfilename}.{rand_name_on_now()}.obj'
print(f'Saving to {path}')
mcubes.export_obj(vertices, triangles, path)
print('Saved')
